package tagging;

import java.util.Iterator;
import java.util.Map;
import java.util.Vector;
import java.util.HashMap;

public class TrainingSet extends Parser {
  HashMap<String, Integer> tagMap_;
  Vector<Tag> tags_;
  SuffixHandler sh_;
  static int totalWordCount_ = 0;
  static int totalTagCount_ = 0;
  
  // Words with count >= UNKNOWN_THRESHOLD will be treated as known.
  public final static int UNKNOWN_THRESHOLD = 2;
  
  public TrainingSet(String fileName) throws Parser.ParserException {
    super(fileName, true);
    
    tagMap_ = new HashMap<String, Integer>();
    tags_ = new Vector<Tag>();
    sh_ = new SuffixHandler();
  }
  
  public void process() throws Parser.ParserException {
    Vector<Parser.ProcessedWord> pwords = new Vector<Parser.ProcessedWord>();
    
    while (parseNextSentence(pwords)) {
      processSentence(pwords);
    }
    
    totalTagCount_ = tags_.size();
    postParseProcess();
  }
  
  // After the parsing is done, runs a post-parse process.
  // For initializing values for good-turing smoothing.
  public void postParseProcess() {
    for (int i = 0; i < tags_.size(); ++i) {
      tags_.get(i).process();
    }
  }
  
  public void processSentence(Vector<Parser.ProcessedWord> pwords) {
    totalWordCount_ += pwords.size();

    for (int i = 0; i < pwords.size() - 1; ++i) {
      Tag t = getTrainingTag(pwords.get(i).getTag());
      t.addWord(pwords.get(i).getWord());
      t.addFollowingTag(pwords.get(i+1).getTag());
      sh_.handleWord(pwords.get(i).getWord(), pwords.get(i).getTag());
    }
    
    Tag t = getTrainingTag(pwords.lastElement().getTag());
    t.addWord(pwords.lastElement().getWord());
    sh_.handleWord(pwords.lastElement().getWord(), pwords.lastElement().getTag());
  }
  
  public Tag getTrainingTag(String tag) {
    Integer tagIndex = tagMap_.get(tag);
    Tag t = null;
    
    if (tagIndex == null) {
      t = new Tag();
      
      tagMap_.put(tag, tags_.size());
      tags_.add(t);
    } else {
      t = tags_.get(tagIndex);
    }

    return t;
  }
  
  public SuffixHandler getSuffixHandler() {
    return sh_;
  }
  
  public Vector<String> getTagStrings() {
    Vector<String> ts = new Vector<String>();
    
    Iterator<Map.Entry<String, Integer>> it = tagMap_.entrySet().iterator();    
    while(it.hasNext()) {
      ts.add(it.next().getKey());
    }
    
    return ts;
  }
  
  public static int getTotalWordCount() {
    return totalWordCount_;
  }
  
  public static int getTotalTagCount() {
    return totalTagCount_;
  }
}
